import onnx
import onnxruntime 
import numpy as np
from onnxruntime.quantization import quantize_static, CalibrationDataReader, QuantType,QuantFormat,CalibrationMethod

ort_version = onnxruntime.__version__
major_version = int(ort_version.split('.')[0])
print(major_version)
# 定义校准数据读取器
class ImageCalibrationDataReader(CalibrationDataReader):
    def __init__(self, calibration_images, input_name,size, model_path=None):
        self.enum_data = None
        self.input_name = input_name
        self.model_size = size
        # 加载校准图像
        self.images = calibration_images
        
        # 初始化模型以获取输入信息
        if model_path:
            session_options = onnxruntime.SessionOptions()
            self.session = onnxruntime.InferenceSession(
                model_path, session_options, providers=["CPUExecutionProvider"])
        
        self.preprocess_flag = True
        self.calib_data = []
        
        # 预处理校准数据
        if self.preprocess_flag:
            self.preprocess()
    
    def get_next(self):
        """返回下一个校准数据批次"""
        if self.enum_data is None:
            self.enum_data = iter(self.calib_data)
        return next(self.enum_data, None)
    
    def preprocess(self):
        """预处理校准图像"""
        for image_path in self.images:
            # 加载和预处理图像
            image = self._load_image(image_path)
            input_tensor = self._preprocess(image)
            
            # 添加到校准数据列表
            self.calib_data.append({self.input_name: input_tensor})
    
    def _load_image(self, image_path):
        """加载图像"""
        import cv2
        return cv2.imread(image_path)
    
    def _preprocess(self, image):
        """图像预处理（根据模型需求调整）"""
        import cv2
        # 调整大小
        image = cv2.resize(image, self.model_size)
        # BGR 转 RGB
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # 归一化
        image = image.astype(np.float32)
    
    #    image = np.transpose(image, [1,2,0])
        return image

# 执行静态量化
def static_quantize_onnx_model(
    input_model_path, 
    output_model_path, 
    calibration_data_reader,
    weight_type=QuantType.QInt8,
    activation_type=QuantType.QUInt8,
    per_channel=True,
    reduce_range=True
):
    """执行 ONNX 模型的静态量化"""
    quantize_static(
        model_input=input_model_path,
        model_output=output_model_path,
        calibration_data_reader=calibration_data_reader,
        weight_type=weight_type,
        activation_type=activation_type,
        per_channel=per_channel,
        calibrate_method=CalibrationMethod.Entropy,
        reduce_range=reduce_range,
        # op_types_to_quantize=["Conv","LeakyRelu","Mul","Add","Sub","Div"],  # 要量化的算子类型 ,"Add","Sub","Div","Transpose","Resize"
        
        op_types_to_quantize=["Conv","LeakyRelu","Mul","Resize","Transpose","Reciprocal"],
         
        extra_options=  {
        "QuantizeBias": True,
        "ActivationSymmetric": False, # 使用激活值量化的非对称模式
        "AddQDQPairToWeight": False,  # 不添加额外的Q/DQ对
        "DedicatedQDQPair": False,    # 不使用专用的Q/DQ对
        "FuseActivation": True        # 融合激活函数
        },
        quant_format=QuantFormat.QOperator
    )
    
    print(f"静态量化完成，模型保存至: {output_model_path}")

# 使用示例
input_model = "cartoon_bg.onnx"
output_model = "cartoon_bg_q8.onnx"
calibration_images = ["test_face.jpg"]  # 校准图像列表
model_size = (320,320)
# 获取模型输入名称
model = onnx.load(input_model)
input_name = model.graph.input[0].name

# 创建校准数据读取器
calibration_reader = ImageCalibrationDataReader(calibration_images, input_name, model_size,input_model)

# 执行静态量化
static_quantize_onnx_model(input_model, output_model, calibration_reader)


from onnxoptimizer import optimize

# 针对 CPU 优化的策略组合
cpu_passes = [
    "fuse_bn_into_conv",  # 减少计算量
    "fuse_pad_into_conv",  # 减少内存访问
    "eliminate_deadend",  # 简化图结构
    "eliminate_unused_initializer",  # 减少内存占用
]
model = onnx.load(output_model)
optimized_for_cpu = optimize(model, cpu_passes)
onnx.save(optimized_for_cpu, "cartoon_bg_q8_s.onnx")




from onnxsim import simplify

# model = onnx.load(output_model)
# # 执行简化
# simplified_model, check = simplify(model,dynamic_input_shape=False)

# # 验证简化结果
# assert check, "简化验证失败!"

# # 保存简化后模型
# onnx.save(simplified_model, 'cartoon_bg_q8_s.onnx')


input_model = "cartoon_h.onnx"
output_model = "cartoon_h_q8.onnx"
calibration_images = ["test_face.jpg"]  # 校准图像列表
model_size = (192,192)
# 获取模型输入名称
model = onnx.load(input_model)
input_name = model.graph.input[0].name

# 创建校准数据读取器
calibration_reader = ImageCalibrationDataReader(calibration_images, input_name, model_size,input_model)

# 执行静态量化
static_quantize_onnx_model(input_model, output_model, calibration_reader)

model = onnx.load(output_model)
# 执行简化
simplified_model, check = simplify(model)

# 验证简化结果
assert check, "简化验证失败!"

# 保存简化后模型
onnx.save(simplified_model, 'cartoon_h_q8_s.onnx')



##python -m onnxsim cartoon_h_q8.onnx cartoon_h_q8_s.onnx
